#pragma once
#include <map>
#include <cstdlib>

// TODO: this should be implemented in terms of a custom self-balancing tree designed for efficient finding an empty key
// and the key part should be  properly made generic

template <typename Key, typename Value>
struct KeyTree
{
	std::size_t size() const;
	bool empty() const;

	Value * insert(Key const&, Value const& );
	Key insert(Value const& );

	bool contains(Key const&) const;

	Value const& at(Key const& ) const;
	Value & at(Key const& );

	void erase(Key const&);

	Key unused_key() const;
private:

	typename std::map<Key, Value>::iterator find(Key const&);
	typename std::map<Key, Value>::const_iterator find(Key const&) const;

	typename std::map<Key, Value> tree_;
};

// implementation

template <typename T>
inline T generate_random() {
	return ::rand();
}

template <typename Key, typename Value>
Value const& KeyTree<Key, Value>::at(const Key &key) const {
	return find(key)->second;
}

template <typename Key, typename Value>
Value & KeyTree<Key, Value>::at(const Key &key) {
	return find(key)->second;
}

template <typename Key, typename Value>
std::size_t KeyTree<Key, Value>::size() const {
	return tree_.size();
}

template <typename Key, typename Value>
bool KeyTree<Key, Value>::empty() const {
	return tree_.empty();
}

template <typename Key, typename Value>
Value * KeyTree<Key, Value>::insert(Key const& k, Value const& v) {
	assert(!contains(k));
	std::pair<typename std::map<Key, Value>::iterator, bool> ret(tree_.insert(std::make_pair(k, v)));
	assert(ret.second == true);

	return &ret.first->second;
}

template <typename Key, typename Value>
Key KeyTree<Key, Value>::unused_key() const
{
	// TODO: this properly...
	while (true) {
		Key k(generate_random<Key>());

		if (!contains(k)) {
			return k;
		}
	}
}

template <typename Key, typename Value>
bool KeyTree<Key, Value>::contains(Key const& k) const {
	return tree_.find(k) != tree_.end();
}

template <typename Key, typename Value>
typename std::map<Key, Value>::iterator KeyTree<Key, Value>::find(Key const& k)
{
	assert(!tree_.empty());

	typename std::multimap<Key, Value>::iterator it(tree_.upper_bound(k-1));
	if (it == tree_.end())
		it = tree_.begin();

	assert (it != tree_.end());

	return it;
}

template <typename Key, typename Value>
typename std::map<Key, Value>::const_iterator KeyTree<Key, Value>::find(Key const& k) const
{
	assert(!tree_.empty());

	typename std::map<Key, Value>::const_iterator it(tree_.upper_bound(k));
	if (it == tree_.end())
		it = tree_.begin();

	assert (it != tree_.end());

	return it;
}

template <typename Key, typename Value>
void KeyTree<Key, Value>::erase(Key const& k)
{
	assert(!tree_.empty());

	tree_.erase(find(k));
}



